"""
05_advanced_queries.py
======================
Advanced Cypher Query Patterns

This module covers:
1. Variable-length path queries
2. Shortest path algorithms
3. Pattern comprehensions
4. Subqueries and CALL
5. UNION and UNWIND
6. Aggregation and grouping
7. Conditional logic (CASE WHEN)
8. Working with collections

Author: Neo4j Tutorial
Date: January 2025
"""

from typing import Dict, List, Any, Optional, Tuple
import sys
import os

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from src.utils.connection import get_driver


class AdvancedQueries:
    """
    Advanced Cypher query patterns for complex graph operations.
    
    Query Complexity Levels:
    
    ┌─────────────────────────────────────────────────────────────────┐
    │                    QUERY COMPLEXITY                             │
    ├─────────────────────────────────────────────────────────────────┤
    │                                                                 │
    │  LEVEL 1: Basic                                                │
    │  ─────────────────                                              │
    │  MATCH (n:Label) RETURN n                                      │
    │                                                                 │
    │  LEVEL 2: Pattern Matching                                     │
    │  ─────────────────────────                                      │
    │  MATCH (a)-[r]->(b) WHERE ... RETURN ...                       │
    │                                                                 │
    │  LEVEL 3: Multi-Hop                                            │
    │  ───────────────────                                            │
    │  MATCH path = (a)-[*1..3]->(b) RETURN path                     │
    │                                                                 │
    │  LEVEL 4: Subqueries                                           │
    │  ────────────────────                                           │
    │  CALL { WITH a MATCH ... RETURN ... }                          │
    │                                                                 │
    │  LEVEL 5: Graph Algorithms                                     │
    │  ─────────────────────────                                      │
    │  shortestPath, allShortestPaths, apoc.path.*                  │
    │                                                                 │
    └─────────────────────────────────────────────────────────────────┘
    """
    
    def __init__(self, driver=None):
        self.driver = driver or get_driver()
    
    # =========================================================
    # VARIABLE-LENGTH PATHS
    # =========================================================
    
    def find_connected_nodes(
        self,
        label: str,
        node_id: str,
        min_hops: int = 1,
        max_hops: int = 3,
        rel_types: List[str] = None
    ) -> List[Dict[str, Any]]:
        """
        Find nodes connected within N hops.
        
        Variable-length path syntax: [*min..max]
        - [*] = any number of hops
        - [*1..3] = 1 to 3 hops
        - [*2] = exactly 2 hops
        
        Args:
            label: Starting node label
            node_id: Starting node ID
            min_hops: Minimum path length
            max_hops: Maximum path length
            rel_types: Optional relationship type filter
            
        Returns:
            Connected nodes with path info
        """
        rel_filter = ""
        if rel_types:
            rel_filter = ":" + "|".join(rel_types)
        
        query = f"""
        MATCH path = (start:{label} {{id: $id}})-[r{rel_filter}*{min_hops}..{max_hops}]-(connected)
        WHERE start <> connected
        RETURN DISTINCT
            labels(connected)[0] as label,
            connected.id as id,
            connected.name as name,
            length(path) as distance,
            [rel in relationships(path) | type(rel)] as relationship_path
        ORDER BY distance, connected.name
        """
        
        with self.driver.session() as session:
            result = session.run(query, id=node_id)
            return [r.data() for r in result]
    
    def find_paths_between(
        self,
        from_label: str,
        from_id: str,
        to_label: str,
        to_id: str,
        max_hops: int = 5
    ) -> List[Dict[str, Any]]:
        """
        Find all paths between two nodes.
        
        Args:
            from_label: Start node label
            from_id: Start node ID
            to_label: End node label
            to_id: End node ID
            max_hops: Maximum path length
            
        Returns:
            All paths with details
        """
        query = f"""
        MATCH path = (a:{from_label} {{id: $from_id}})-[*1..{max_hops}]-(b:{to_label} {{id: $to_id}})
        RETURN 
            [node in nodes(path) | node.name] as node_names,
            [rel in relationships(path) | type(rel)] as relationships,
            length(path) as path_length
        ORDER BY path_length
        LIMIT 10
        """
        
        with self.driver.session() as session:
            result = session.run(query, from_id=from_id, to_id=to_id)
            return [r.data() for r in result]
    
    # =========================================================
    # SHORTEST PATH ALGORITHMS
    # =========================================================
    
    def shortest_path(
        self,
        from_label: str,
        from_id: str,
        to_label: str,
        to_id: str,
        rel_types: List[str] = None
    ) -> Optional[Dict[str, Any]]:
        """
        Find the shortest path between two nodes.
        
        Uses Cypher's built-in shortestPath function.
        
        Args:
            from_label: Start node label
            from_id: Start node ID
            to_label: End node label
            to_id: End node ID
            rel_types: Optional relationship type filter
            
        Returns:
            Shortest path details or None
        """
        rel_filter = ""
        if rel_types:
            rel_filter = ":" + "|".join(rel_types)
        
        query = f"""
        MATCH (a:{from_label} {{id: $from_id}}), (b:{to_label} {{id: $to_id}})
        MATCH path = shortestPath((a)-[{rel_filter}*]-(b))
        RETURN 
            [node in nodes(path) | {{label: labels(node)[0], name: node.name}}] as nodes,
            [rel in relationships(path) | type(rel)] as relationships,
            length(path) as path_length
        """
        
        with self.driver.session() as session:
            result = session.run(query, from_id=from_id, to_id=to_id)
            record = result.single()
            return record.data() if record else None
    
    def all_shortest_paths(
        self,
        from_label: str,
        from_id: str,
        to_label: str,
        to_id: str,
        max_paths: int = 5
    ) -> List[Dict[str, Any]]:
        """
        Find all shortest paths between two nodes.
        
        If multiple paths have the same length, returns all of them.
        """
        query = f"""
        MATCH (a:{from_label} {{id: $from_id}}), (b:{to_label} {{id: $to_id}})
        MATCH path = allShortestPaths((a)-[*]-(b))
        WITH path LIMIT $max_paths
        RETURN 
            [node in nodes(path) | node.name] as node_names,
            [rel in relationships(path) | type(rel)] as relationships,
            length(path) as path_length
        """
        
        with self.driver.session() as session:
            result = session.run(query, from_id=from_id, to_id=to_id, max_paths=max_paths)
            return [r.data() for r in result]
    
    # =========================================================
    # PATTERN COMPREHENSIONS
    # =========================================================
    
    def get_node_with_related(
        self,
        label: str,
        node_id: str
    ) -> Dict[str, Any]:
        """
        Get node with all related entities using pattern comprehension.
        
        Pattern comprehension syntax:
        [(pattern) WHERE condition | expression]
        
        This is like a list comprehension for graph patterns.
        """
        query = f"""
        MATCH (n:{label} {{id: $id}})
        RETURN 
            n.name as name,
            n.id as id,
            [(n)-[:DIAGNOSED_WITH]->(d:Disease) | d.name] as diseases,
            [(n)-[:TAKES]->(m:Medication) | {{
                name: m.name,
                class: m.drugClass
            }}] as medications,
            [(n)-[:TREATED_BY]->(doc:Doctor) | doc.name] as doctors
        """
        
        with self.driver.session() as session:
            result = session.run(query, id=node_id)
            record = result.single()
            return record.data() if record else None
    
    # =========================================================
    # SUBQUERIES
    # =========================================================
    
    def correlated_subquery(
        self,
        label: str
    ) -> List[Dict[str, Any]]:
        """
        Use CALL subquery for complex aggregations.
        
        Subqueries allow you to:
        - Perform aggregations per row
        - Use LIMIT within a subquery
        - Execute complex logic per matched node
        """
        query = f"""
        MATCH (n:{label})
        CALL {{
            WITH n
            OPTIONAL MATCH (n)-[r:DIAGNOSED_WITH]->(d:Disease)
            RETURN count(d) as disease_count
        }}
        CALL {{
            WITH n
            OPTIONAL MATCH (n)-[r:TAKES]->(m:Medication)
            WHERE r.status = 'active'
            RETURN count(m) as medication_count
        }}
        RETURN 
            n.id as id,
            n.name as name,
            disease_count,
            medication_count
        ORDER BY disease_count DESC, medication_count DESC
        LIMIT 20
        """
        
        with self.driver.session() as session:
            result = session.run(query)
            return [r.data() for r in result]
    
    def exists_subquery(
        self,
        label: str,
        related_label: str,
        rel_type: str
    ) -> List[Dict[str, Any]]:
        """
        Use EXISTS subquery for filtering.
        
        EXISTS is efficient for checking relationship existence
        without needing to return the matched data.
        """
        query = f"""
        MATCH (n:{label})
        WHERE EXISTS {{
            MATCH (n)-[:{rel_type}]->(:{related_label})
        }}
        RETURN n.id as id, n.name as name
        ORDER BY n.name
        """
        
        with self.driver.session() as session:
            result = session.run(query)
            return [r.data() for r in result]
    
    # =========================================================
    # UNION AND UNWIND
    # =========================================================
    
    def union_query(
        self,
        patient_id: str
    ) -> List[Dict[str, Any]]:
        """
        Combine results from multiple queries with UNION.
        
        UNION: Removes duplicates
        UNION ALL: Keeps duplicates (faster)
        """
        query = """
        MATCH (p:Patient {id: $id})-[:DIAGNOSED_WITH]->(d:Disease)
        RETURN 'Disease' as type, d.name as name, d.id as code
        
        UNION ALL
        
        MATCH (p:Patient {id: $id})-[:TAKES]->(m:Medication)
        RETURN 'Medication' as type, m.name as name, m.id as code
        
        UNION ALL
        
        MATCH (p:Patient {id: $id})-[:ALLERGIC_TO]->(a)
        RETURN 'Allergy' as type, a.name as name, a.id as code
        """
        
        with self.driver.session() as session:
            result = session.run(query, id=patient_id)
            return [r.data() for r in result]
    
    def unwind_and_process(
        self,
        items: List[Dict[str, Any]]
    ) -> List[Dict[str, Any]]:
        """
        Process a list of items using UNWIND.
        
        UNWIND expands a list into rows, allowing you to
        process each item individually.
        """
        query = """
        UNWIND $items AS item
        MERGE (n:ProcessedItem {id: item.id})
        SET n.name = item.name,
            n.value = item.value,
            n.processedAt = datetime()
        RETURN n.id as id, n.name as name
        """
        
        with self.driver.session() as session:
            result = session.run(query, items=items)
            return [r.data() for r in result]
    
    # =========================================================
    # AGGREGATION AND GROUPING
    # =========================================================
    
    def group_and_aggregate(
        self,
        label: str,
        group_by: str
    ) -> List[Dict[str, Any]]:
        """
        Group by a property and calculate aggregations.
        
        Aggregation functions:
        - count(), count(DISTINCT x)
        - sum(), avg(), min(), max()
        - collect() - creates a list
        - stDev() - standard deviation
        """
        query = f"""
        MATCH (n:{label})
        WHERE n.{group_by} IS NOT NULL
        WITH n.{group_by} as group_value, n
        RETURN 
            group_value,
            count(n) as total,
            avg(n.age) as avg_age,
            min(n.age) as min_age,
            max(n.age) as max_age,
            collect(n.name)[0..3] as sample_names
        ORDER BY total DESC
        """
        
        with self.driver.session() as session:
            result = session.run(query)
            return [r.data() for r in result]
    
    def pivot_query(
        self,
        patient_label: str = "Patient"
    ) -> List[Dict[str, Any]]:
        """
        Create a pivot-like result showing patients with their conditions.
        """
        query = f"""
        MATCH (p:{patient_label})
        OPTIONAL MATCH (p)-[d:DIAGNOSED_WITH]->(disease:Disease)
        OPTIONAL MATCH (p)-[t:TAKES]->(med:Medication)
        WITH p, 
             collect(DISTINCT disease.name) as diseases,
             collect(DISTINCT med.name) as medications
        RETURN 
            p.id as patient_id,
            p.name as patient_name,
            size(diseases) as disease_count,
            diseases,
            size(medications) as medication_count,
            medications
        ORDER BY disease_count DESC
        LIMIT 20
        """
        
        with self.driver.session() as session:
            result = session.run(query)
            return [r.data() for r in result]
    
    # =========================================================
    # CONDITIONAL LOGIC
    # =========================================================
    
    def case_when_query(
        self,
        label: str
    ) -> List[Dict[str, Any]]:
        """
        Use CASE WHEN for conditional logic.
        
        Syntax:
        CASE 
            WHEN condition1 THEN result1
            WHEN condition2 THEN result2
            ELSE default
        END
        """
        query = f"""
        MATCH (n:{label})
        RETURN 
            n.id as id,
            n.name as name,
            n.age as age,
            CASE 
                WHEN n.age < 18 THEN 'Minor'
                WHEN n.age < 40 THEN 'Young Adult'
                WHEN n.age < 65 THEN 'Middle Age'
                ELSE 'Senior'
            END as age_group,
            CASE n.gender
                WHEN 'Male' THEN 'M'
                WHEN 'Female' THEN 'F'
                ELSE 'O'
            END as gender_code
        ORDER BY n.age
        """
        
        with self.driver.session() as session:
            result = session.run(query)
            return [r.data() for r in result]
    
    # =========================================================
    # COLLECTION OPERATIONS
    # =========================================================
    
    def list_operations(
        self,
        label: str,
        node_id: str
    ) -> Dict[str, Any]:
        """
        Demonstrate list operations in Cypher.
        
        List functions:
        - head(), last(), tail()
        - size(), range()
        - reduce(), [x IN list WHERE pred | expr]
        """
        query = f"""
        MATCH (n:{label} {{id: $id}})-[:DIAGNOSED_WITH]->(d:Disease)
        WITH n, collect(d.name) as diseases
        RETURN 
            n.name as patient,
            diseases,
            size(diseases) as count,
            head(diseases) as first_disease,
            last(diseases) as last_disease,
            tail(diseases) as rest_of_diseases,
            [d IN diseases WHERE size(d) > 10] as long_names,
            [d IN diseases | toUpper(d)] as uppercase_names
        """
        
        with self.driver.session() as session:
            result = session.run(query, id=node_id)
            record = result.single()
            return record.data() if record else None
    
    def reduce_operation(
        self
    ) -> List[Dict[str, Any]]:
        """
        Use reduce() for accumulating values across a collection.
        
        reduce(accumulator = initial, x IN list | expression)
        """
        query = """
        MATCH (p:Patient)-[t:TAKES]->(m:Medication)
        WHERE t.status = 'active'
        WITH p, collect(m.name) as medications
        RETURN 
            p.name as patient,
            medications,
            reduce(s = '', med IN medications | s + med + ', ') as med_string,
            reduce(count = 0, med IN medications | count + 1) as manual_count
        LIMIT 10
        """
        
        with self.driver.session() as session:
            result = session.run(query)
            return [r.data() for r in result]


# ============================================================
# DEMONSTRATION
# ============================================================

if __name__ == "__main__":
    print("=" * 60)
    print("Advanced Cypher Queries - Demonstration")
    print("=" * 60)
    
    try:
        queries = AdvancedQueries()
        
        # Variable-length paths
        print("\n📌 Variable-Length Paths")
        print("-" * 40)
        print("   Finding nodes within 1-3 hops...")
        
        # Subqueries
        print("\n📌 Subquery Examples")
        print("-" * 40)
        print("   Correlated subqueries for per-row aggregation...")
        
        # Aggregations
        print("\n📌 Aggregation Patterns")
        print("-" * 40)
        print("   GROUP BY with collect(), count(), avg()...")
        
        # CASE WHEN
        print("\n📌 Conditional Logic")
        print("-" * 40)
        print("   CASE WHEN for age group classification...")
        
        print("\n" + "=" * 60)
        print("✅ Advanced queries demonstration complete!")
        print("=" * 60)
        
    except Exception as e:
        print(f"Error: {e}")
